package problems.algorithm.prefix_tree;

import java.util.HashMap;

/**
 * 6183. 字符串的前缀分数和
 * <p>https://leetcode.cn/problems/sum-of-prefix-scores-of-strings/</p>
 */
public class T6183 {
    public int[] sumPrefixScores(String[] words) {
        // 建立前缀树
        TrieNode root = new TrieNode();
        for (String word : words) {
            insert(root, word);
        }

        // 建立 HashMap 进行加速
        HashMap<String, Integer> cntMap = new HashMap<>();
        int n = words.length;
        int[] ret = new int[n];
//        String key;

        for (int i = 0; i < n; ++i) {
            if (cntMap.containsKey(words[i])) {
                ret[i] = cntMap.get(words[i]);
                continue;
            }
            ret[i] = searchPrefix(root, words[i]);
            cntMap.put(words[i], ret[i]);
        }

        return ret;

    }

    // 以小写字母组成的单词组为例
    public static class TrieNode {
        // 前缀的个数
        public int prefix;
        TrieNode[] nextNode = new TrieNode[26];

        public TrieNode() {
            prefix = 0;
        }
    }

    // 插入一个单词
    public void insert(TrieNode root, String word) {
        if (root == null || word == null || word.length() < 1) {
            return;
        }

        char[] chs = word.toCharArray();
        int n = chs.length;
        int j;

        for (int i = 0; i < n; ++i) {
            j = chs[i] - 'a';
            if (root.nextNode[j] == null) {
                root.nextNode[j] = new TrieNode();
            }
            root = root.nextNode[j];
            ++root.prefix;
        }
    }

    // 查询以str为前缀的单词数量
    public int searchPrefix(TrieNode root, String str) {
        if (root == null || str == null || str.length() < 1) return 0;
        char[] chs = str.toCharArray();
        int ret = 0;

        for (char c : chs) {
            if (root.nextNode[c - 'a'] == null) return 0;
            root = root.nextNode[c - 'a'];
            ret += root.prefix;
        }
        return ret;
    }
}
